import random

import cv2
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image, ImageDraw


def generate_touch_captcha(size=(300, 300), num_objects=5):
    """生成点触验证码：随机绘制图形并标记目标位置"""
    img = Image.new("RGB", size, (255, 255, 255))
    draw = ImageDraw.Draw(img)
    target_info = ("circle", (255, 0, 0))  # (形状, RGB颜色) 目标为红色圆形
    objects = []  # 存储目标位置

    # 强制生成1个目标红色圆形（确保至少有1个目标）
    x_target, y_target = (
        random.randint(50, size[0] - 50),
        random.randint(50, size[1] - 50),
    )
    draw.ellipse(
        [(x_target - 20, y_target - 20), (x_target + 20, y_target + 20)],
        fill=target_info[1],
    )
    objects.append((x_target, y_target))

    # 生成剩余随机图形
    for _ in range(num_objects - 1):
        x, y = random.randint(50, size[0] - 50), random.randint(50, size[1] - 50)
        shape = random.choice(["circle", "square"])
        color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))

        if shape == "circle":
            draw.ellipse([(x - 20, y - 20), (x + 20, y + 20)], fill=color)
        else:
            draw.rectangle([(x - 20, y - 20), (x + 20, y + 20)], fill=color)

    # 转换颜色格式
    img_bgr = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
    return (
        Image.fromarray(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)),
        objects,
        target_info,
    )


def solve_touch_captcha(image, target_rgb):
    """通过颜色检测破解点触验证码（修正BGR颜色空间问题）"""
    # 将Pillow图像转换为OpenCV的BGR格式
    img_bgr = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)

    # 将目标RGB颜色转换为BGR格式
    target_bgr = (target_rgb[2], target_rgb[1], target_rgb[0])

    # 扩大颜色过滤范围（原±15改为±30）
    lower = np.array([target_bgr[0] - 30, target_bgr[1] - 30, target_bgr[2] - 30])
    upper = np.array([target_bgr[0] + 30, target_bgr[1] + 30, target_bgr[2] + 30])
    mask = cv2.inRange(img_bgr, lower, upper)

    # 轮廓检测（增加轮廓面积过滤，避免误检小噪点）
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    centers = []
    for cnt in contours:
        area = cv2.contourArea(cnt)
        if area > 100:  # 只保留面积大于100的轮廓（目标图形面积约为π*20²≈1256）
            M = cv2.moments(cnt)
            if M["m00"] != 0:
                cx, cy = int(M["m10"] / M["m00"]), int(M["m01"] / M["m00"])
                centers.append((cx, cy))
    return centers


def main():
    # 生成并破解
    captcha_img, true_pos, (target_shape, target_color) = generate_touch_captcha()
    solved_pos = solve_touch_captcha(captcha_img, target_color)

    # 可视化对比
    plt.figure(figsize=(12, 6))
    # 原图+正确位置（绿叉）
    plt.subplot(1, 2, 1)
    plt.imshow(captcha_img)
    plt.title(f"Original (Target: {target_shape} {target_color})")
    for x, y in true_pos:
        plt.scatter(x, y, c="green", marker="X", s=150)

    # 破解结果（红圈）
    plt.subplot(1, 2, 2)
    plt.imshow(captcha_img)
    plt.title(f"Solved (Found {len(solved_pos)} positions)")
    for x, y in solved_pos:
        plt.scatter(x, y, c="yellow", marker="v", s=150)

    plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    main()
